# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

layer_configs = [
    # Unit1 (2)
    (32, 2, True),
    (63, 3, True),
    # Unit2 (3)
    (128, 3, False),
    (64, 1, False),
    (128, 3, True),
    # Unit3 (3)
    (256, 3 False),
    (128, 1, False),
    (256, 3, True),
    # Unit4 (5)
    (512, 3, False),
    (256, 1, False),
    (512, 3, False),
    (256, 1, False),
    (512, 3, True),
    # Unit5 (5)
    (1024, 3, False),
    (512, 3, False),
    (1024, 3, False),
    (512, 1, False),
    (1024, 3, False),
]

class conv_block(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, pool, stride=1):
        super(conv_block, self).__init__()
        pad = 1 if kernel_size == 3 else 0
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size,
                              stride=stride, padding=pad, bias=False)
        self.bn = nn.BatchNorm2d(out_channels)
        self.act = nn.LeakyReLU(0.1)
        self.pool= pool

    def forward(self, x):
        out = self.conv(x)
        out = self.bn(out)
        out = self.act(out)
        if self.pool:
            out = F.max_pool(out, kernel_size=2, strider=2)
        return out

class darknet19(nn.Module):
    def __init__(self, cls_num=1000):
        super(darknet19, self).__init__()
        self.cls_num = cls_num
        self.feature = self.make_layers(3, layer_configs)

    def make_layers(self, in_channels, cfg):
        layers = []
        for out_channels, kernel_size, pool in cfg:
            layers.append(conv_block(in_channels, out_channels, kernel_size,
                                     pool))
            in_channels = out_channels
        return nn.Sequential(*layers)

    def load_weight(self, weight_file):
        print("Load pretrained models!")

        fp = open(weight_file, 'rb')
        header = np.fromfile(fp, count=4, dtype=np.int32)
        header = torch.from_numpy(header)
        buf = np.fromfile(fp, dtype=np.float32)
        start = 0
        for idx, m in enumerate(self.feature.modules()):
            if isinstance(m, nn.Conv2d):
                conv = m
            if isinstance(m, nn.BatchNorm2d):
                bn = m
                start = load_conv_bn(buf, start, conv, bn)
        
        assert start == buf.shape[0]

    def forward(self, x):
        output = self.feature(x)
        return output

class YOLO(nn.Module):
    def __init__(self, cls_num, bbox_num = 2, scale_size=7, conv_model=True,
                 pretrained=None):
        super(YOLO, self).__init__()
        self.cls_num = cls_num
        self.feature = darknet19()
        self.conv_model = conv_model
        if pretrained is not None:
            self.feature.load_weight(pretrained)

        self.scale_size = scale_size
        self.bbox_num = bbox_num
        self.last_output = (5 * self.bbox_num + self.cls_num)

        cx = torch.linspace(0.5 / scale_size, (scale_size - 0.5) / scale_size, steps=scale_size.view(-1, scale_size).repeat((scale_size, 1)).view(scale_size, scale_size, -1))

